from easycore.common.registry import Registry
import os
import matplotlib.pyplot as plt

ACTION_REGISTRY = Registry('action')


class DrawFigures:
    def __init__(self):
        self.name = 'draw figures'

    def __call__(self, evaluate_dict_list, output_dir):
        threshold_list = [evaluate_dict['threshold'] for evaluate_dict in evaluate_dict_list]
        proposal_count_list = [evaluate_dict['average_proposal_count'] for evaluate_dict in evaluate_dict_list]
        for metric_type in ['macro', 'micro']:
            for strategy in ['iou_cover', 'proposals_cover', 'gt_boxes_cover']:
                precision_list = [evaluate_dict[strategy]['{}_precision'.format(metric_type)] for evaluate_dict in evaluate_dict_list]
                recall_list = [evaluate_dict[strategy]['{}_recall'.format(metric_type)] for evaluate_dict in evaluate_dict_list]
                
                fig = plt.figure()
                plt.plot(threshold_list, precision_list, label='precision')
                plt.plot(threshold_list, recall_list, label='recall')
                plt.legend()
                plt.title("{}-{}".format(strategy, metric_type))
                plt.xlim(0,1)
                plt.ylim(-0.01,1.01)
                plt.xlabel('threshold')
                plt.ylabel('value')
                plt.savefig(os.path.join(output_dir, '{}-{}.png'.format(strategy, metric_type)), dpi=240)

        fig = plt.figure()
        plt.plot(threshold_list, proposal_count_list)
        plt.xlabel('threshold')
        plt.ylabel('average number of proposals')
        plt.title('proposals count')
        plt.savefig(os.path.join(output_dir, 'proposal_count.png'), dpi=240)

action = DrawFigures()
ACTION_REGISTRY.register(action.name, action)




class DrawFiguresPlus:
    def __init__(self):
        self.name = 'draw figures plus'

    def __call__(self, evaluate_dict_list, output_dir):
        threshold_list = [evaluate_dict['threshold'] for evaluate_dict in evaluate_dict_list]
        proposal_count_list = [evaluate_dict['average_proposal_count'] for evaluate_dict in evaluate_dict_list]
        for metric_type in ['macro', 'micro']:
            for strategy in ['iou_cover', 'proposals_cover', 'gt_boxes_cover']:
                precision_list = [evaluate_dict[strategy]['{}_precision'.format(metric_type)] for evaluate_dict in evaluate_dict_list]
                recall_list = [evaluate_dict[strategy]['{}_recall'.format(metric_type)] for evaluate_dict in evaluate_dict_list]
                
                fig = plt.figure()
                fig, ax1 = plt.subplots()
                ax2 = ax1.twinx()
                ax1.plot(threshold_list, precision_list, label='precision')
                ax1.plot(threshold_list, recall_list, label='recall')
                ax1.legend(loc='lower left')
                ax1.grid()
                plt.title("{}-{}".format(strategy, metric_type))
                ax1.set_xlim(0,1)
                ax1.set_ylim(0,1.01)
                ax1.set_xlabel('threshold')
                ax1.set_ylabel('precision or recall')

                ax2.plot(threshold_list, proposal_count_list, color='g', label='proposal count')
                ax2.set_ylabel('average number of proposals')
                ax2.legend(loc='upper right')
                ax2.set_ylim(ymin=0)

                plt.savefig(os.path.join(output_dir, '{}-{}-plus.png'.format(strategy, metric_type)), dpi=240)

action = DrawFiguresPlus()
ACTION_REGISTRY.register(action.name, action)